#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Camera Node Manager for Muto Robot

Provides persistent camera node management to avoid frequent node creation/destruction
during have_a_look operations.
"""

import time
import threading
import uuid
from typing import Optional, Tuple, Any, Dict

import rclpy
from rclpy.node import Node
from sensor_msgs.msg import Image
from cv_bridge import CvBridge
import cv2
import numpy as np


class PersistentImageCaptureNode(Node):
    """持久化的图像捕获节点"""
    
    def __init__(self, image_topic: str = "/camera/color/image_raw"):
        """
        初始化持久化图像捕获节点
        
        Args:
            image_topic: 图像话题名称
        """
        node_id = str(uuid.uuid4())[:8]
        super().__init__(f'persistent_image_capture_{node_id}')
        
        self.image_topic = image_topic
        self.bridge = CvBridge()
        self.latest_image = None
        self.image_timestamp = None
        self.image_lock = threading.Lock()
        self.is_active = True
        
        # 创建图像订阅器
        self.subscription = self.create_subscription(
            Image,
            self.image_topic,
            self.image_callback,
            10
        )
        
        self.get_logger().info(f'Persistent image capture node initialized for topic: {self.image_topic}')
    
    def image_callback(self, msg):
        """图像回调函数"""
        try:
            # 将ROS图像消息转换为OpenCV图像
            cv_image = self.bridge.imgmsg_to_cv2(msg, "bgr8")
            
            with self.image_lock:
                self.latest_image = cv_image.copy()
                self.image_timestamp = time.time()
                
        except Exception as e:
            self.get_logger().error(f'Error in image callback: {e}')
    
    def get_latest_image(self, max_age_seconds: float = 2.0) -> Tuple[bool, Optional[np.ndarray], str]:
        """
        获取最新的图像
        
        Args:
            max_age_seconds: 图像最大年龄（秒）
            
        Returns:
            Tuple[bool, Optional[np.ndarray], str]: (成功标志, 图像数据, 消息)
        """
        with self.image_lock:
            if self.latest_image is None:
                return False, None, "No image received yet"
            
            if self.image_timestamp is None:
                return False, None, "No image timestamp available"
            
            age = time.time() - self.image_timestamp
            if age > max_age_seconds:
                return False, None, f"Image too old: {age:.2f}s > {max_age_seconds}s"
            
            return True, self.latest_image.copy(), f"Image captured successfully (age: {age:.2f}s)"
    
    def is_receiving_images(self, check_duration: float = 1.0) -> bool:
        """
        检查是否正在接收图像
        
        Args:
            check_duration: 检查持续时间（秒）
            
        Returns:
            bool: 是否正在接收图像
        """
        initial_timestamp = self.image_timestamp
        time.sleep(check_duration)
        return self.image_timestamp != initial_timestamp and self.image_timestamp is not None
    
    def shutdown(self):
        """关闭节点"""
        self.is_active = False
        self.get_logger().info('Persistent image capture node shutting down')


class CameraNodeManager:
    """摄像头节点管理器"""
    
    def __init__(self, logger=None):
        """
        初始化摄像头节点管理器
        
        Args:
            logger: 日志记录器
        """
        self.logger = logger
        self.image_capture_node = None
        self.node_lock = threading.Lock()
        self.is_initialized = False
        self.consecutive_failures = 0
        self.next_recover_ts = 0.0
        self.max_backoff_seconds = 30.0
        
    def initialize_camera_node(self, image_topic: str = "/camera/color/image_raw") -> Tuple[bool, str]:
        """
        初始化摄像头节点
        
        Args:
            image_topic: 图像话题名称
            
        Returns:
            Tuple[bool, str]: (成功标志, 消息)
        """
        with self.node_lock:
            try:
                # 检查ROS2上下文
                if not rclpy.ok():
                    return False, "ROS2 context not initialized"
                
                # 如果节点已存在，先清理
                if self.image_capture_node is not None:
                    self._cleanup_node()
                
                # 创建新的持久化图像捕获节点
                self.image_capture_node = PersistentImageCaptureNode(image_topic)
                self.is_initialized = True
                
                if self.logger:
                    self.logger.info(f'Camera node initialized successfully for topic: {image_topic}')
                
                return True, "Camera node initialized successfully"
                
            except Exception as e:
                error_msg = f"Failed to initialize camera node: {e}"
                if self.logger:
                    self.logger.error(error_msg)
                return False, error_msg
    
    def capture_image_fast(self, save_path: str, timeout_s: float = 5.0, 
                          image_format: str = "jpg") -> Tuple[bool, str, Dict[str, Any]]:
        """
        快速图像捕获（使用持久化节点）
        
        Args:
            save_path: 保存路径
            timeout_s: 超时时间
            image_format: 图像格式
            
        Returns:
            Tuple[bool, str, Dict[str, Any]]: (成功标志, 消息, 数据)
        """
        with self.node_lock:
            try:
                # 检查节点是否已初始化
                if not self.is_initialized or self.image_capture_node is None:
                    return False, "Camera node not initialized", {"error": "node_not_initialized"}
                
                # 等待图像数据
                start_time = time.time()
                while time.time() - start_time < timeout_s:
                    success, image, message = self.image_capture_node.get_latest_image()
                    if success:
                        # 保存图像
                        try:
                            cv2.imwrite(save_path, image)
                            
                            # 验证文件是否成功保存
                            import os
                            if os.path.exists(save_path) and os.path.getsize(save_path) > 0:
                                if self.logger:
                                    self.logger.info(f'Image saved successfully: {save_path}')
                                
                                return True, f"Image captured and saved: {save_path}", {
                                    "save_path": save_path,
                                    "image_shape": image.shape,
                                    "file_size": os.path.getsize(save_path),
                                    "capture_method": "persistent_node"
                                }
                            else:
                                return False, "Failed to save image file", {"error": "file_save_failed"}
                                
                        except Exception as save_error:
                            return False, f"Error saving image: {save_error}", {"error": "save_exception"}
                    
                    # 短暂等待
                    time.sleep(0.1)
                
                # 超时
                return False, f"Image capture timeout after {timeout_s}s", {"error": "timeout"}
                
            except Exception as e:
                error_msg = f"Error in fast image capture: {e}"
                if self.logger:
                    self.logger.error(error_msg)
                return False, error_msg, {"error": "capture_exception"}
    
    def check_camera_node_health(self) -> Tuple[bool, str, Dict[str, Any]]:
        """
        检查摄像头节点健康状态
        
        Returns:
            Tuple[bool, str, Dict[str, Any]]: (健康标志, 消息, 状态数据)
        """
        with self.node_lock:
            try:
                if not self.is_initialized or self.image_capture_node is None:
                    return False, "Camera node not initialized", {"status": "not_initialized"}
                
                # 检查节点是否还活跃
                if not self.image_capture_node.is_active:
                    return False, "Camera node is not active", {"status": "inactive"}
                
                # 检查是否正在接收图像
                receiving_images = self.image_capture_node.is_receiving_images(0.5)
                
                if receiving_images:
                    self.consecutive_failures = 0
                    self.next_recover_ts = 0.0
                    return True, "Camera node is healthy and receiving images", {
                        "status": "healthy",
                        "receiving_images": True,
                        "latest_image_age": time.time() - (self.image_capture_node.image_timestamp or 0)
                    }
                else:
                    self.consecutive_failures += 1
                    backoff = min(self.max_backoff_seconds, 2 ** self.consecutive_failures)
                    self.next_recover_ts = time.time() + backoff
                    return False, "Camera node not receiving images", {
                        "status": "no_images",
                        "receiving_images": False,
                        "consecutive_failures": self.consecutive_failures,
                        "next_recovery_in": backoff
                    }
                    
            except Exception as e:
                error_msg = f"Error checking camera node health: {e}"
                if self.logger:
                    self.logger.error(error_msg)
                return False, error_msg, {"status": "error", "error": str(e)}
    
    def _cleanup_node(self):
        """清理现有节点"""
        try:
            if self.image_capture_node is not None:
                self.image_capture_node.shutdown()
                self.image_capture_node.destroy_node()
                self.image_capture_node = None
                self.is_initialized = False
                
                if self.logger:
                    self.logger.info('Camera node cleaned up')
                    
        except Exception as e:
            if self.logger:
                self.logger.error(f'Error cleaning up camera node: {e}')
    
    def shutdown(self):
        """关闭管理器"""
        with self.node_lock:
            self._cleanup_node()
            if self.logger:
                self.logger.info('Camera node manager shutdown')

    def attempt_recovery(self, image_topic: str = "/camera/color/image_raw") -> Tuple[bool, str]:
        """在指数退避到期后尝试重建摄像头节点并恢复订阅"""
        with self.node_lock:
            now = time.time()
            if self.next_recover_ts and now < self.next_recover_ts:
                remaining = self.next_recover_ts - now
                return False, f"Backoff in effect, try again in {remaining:.2f}s"
            # 清理并重新初始化
            self._cleanup_node()
            ok, msg = self.initialize_camera_node(image_topic)
            if ok:
                self.consecutive_failures = 0
                self.next_recover_ts = 0.0
                return True, "Recovery successful"
            else:
                # 加大退避时间
                self.consecutive_failures += 1
                backoff = min(self.max_backoff_seconds, 2 ** self.consecutive_failures)
                self.next_recover_ts = time.time() + backoff
                return False, f"Recovery failed: {msg}. Next attempt in {backoff:.2f}s"